/*****************************************************************************
*
* AVRPROG compatible boot-loader
* Version  : 0.80beta3 (May 2006)
* Compiler : avr-gcc 3.4.6 / avr-libc 1.4.4
* size     : depends on features and startup ( minmal features < 512 words)
* by       : Martin Thomas, Kaiserslautern, Germany
*            eversmith@heizung-thomas.de
*            Additional code and improvements contributed by:
*           - Uwe Bonnes
*           - Bjoern Riemer
*           - Olaf Rempel
*
* License  : Copyright (c) 2006 Martin Thomas
*            Free to use. You have to mention the copyright
*            owners in source-code and documentation of derived
*            work. No warranty!
*
* Tested with ATmega8, ATmega16, ATmega32, ATmega128, AT90CAN128
*
* - Initial versions have been based on the Butterfly bootloader-code
*   by Atmel Corporation (Authors: BBrandal, PKastnes, ARodland, LHM)
*
****************************************************************************
*
*  See the makefile for information how to adapt the linker-settings to
*  the selected Boot Size (BOOTSIZE=xxxx), the AVR clock-frequency and the
*  MCU-type in
*
*  With BOOT_SIMPLE, minimal features and discarded int-vectors
*  this bootloader has
*  0x2CA - atmega8
*  0x308 - atmega16
*  0x322 - atmega169
*  0x308 - atmega32
*  0x34C - atmega128
*  0x352 - at90can128
*  bytes size and should fit into a 512 word (0x400 byte) bootloader-section.
*  The above values depend on the compiler-version and may differ from
*  your results.
*
****************************************************************************/
/*
	TODOs:
	- check lock-bits set
	- __bad_interrupt still linked even with modified 
	  linker-scripts which needs a default-handler,
	  "wasted": 3 words for AVR5 (>8kB), 2 words for AVR4
*/
// tabsize: 4

// Define boot speeds here
//#define SPEED 16
//#define SPEED 12
#define SPEED 12
#define HWVER 31

#if SPEED == 16
#define F_CPU (16000000)
#define BAUDRATE 500000
#endif

#if SPEED == 12
#define F_CPU (12000000)
#define BAUDRATE 375000
#endif

#if SPEED == 8
#define F_CPU (8000000)
#define BAUDRATE 250000
#endif

// Bootup byte definitions
#if HWVER == 4

	#define DDR_SCK 	DDRE
	#define PIN_SCK     PINE
	#define SCKNUM		PINE7
	#define PORT_SCK	PORTE

	#if SPEED == 16
		#define BOOTBYTE 0x40
	#elif SPEED == 12
		#define BOOTBYTE 0x42
	#else
		#error Unknown speed
	#endif

#elif HWVER == 31

	#define DDR_SCK 	DDRE
	#define PIN_SCK     PINE
	#define SCKNUM		PINE7
	#define PORT_SCK	PORTE

	#if SPEED == 8
		#define BOOTBYTE 0x76
	#endif
	#if SPEED == 12
		#define BOOTBYTE 0x78
	#endif
	#if SPEED == 16
		#define BOOTBYTE 0x7A
	#endif
#else

	#define DDR_SCK 	DDRB
	#define PIN_SCK     PINB
	#define SCKNUM		PINB1
	#define PORT_SCK	PORTB

	#if SPEED == 8
		#define BOOTBYTE 0x71
	#endif
	#if SPEED == 12
		#define BOOTBYTE 0x73
	#endif
	#if SPEED == 16
		#define BOOTBYTE 0x75
	#endif

#endif

/* use "Double Speed Operation" */
//#define UART_DOUBLESPEED

/* use second UART on mega128 / can128 */
//#define UART_USE_SECOND

/* Device-Type:
   For AVRProg the BOOT-option is prefered 
   which is the "correct" value for a bootloader.
   avrdude may only detect the part-code for ISP */
#define DEVTYPE     DEVTYPE_BOOT
//#define DEVTYPE     DEVTYPE_ISP

/*
 * Pin "STARTPIN" on port "STARTPORT" in this port has to grounded
 * (active low) to start the bootloader
 */




/*
 * Select startup-mode
 * SIMPLE-Mode - Jump to bootloader main BL-loop if key is
 *   pressed (Pin grounded) "during" reset or jump to the
 *   application if the pin is not grounded. The internal
 *   pull-up resistor is enabled during the startup and
 *   gets disabled before the application is started.
 * POWERSAVE-Mode - Startup is separated in two loops
 *   which makes power-saving a little easier if no firmware
 *   is on the chip. Needs more memory
 * BOOTICE-Mode - to flash the JTAGICE upgrade.ebn file.
 *   No startup-sequence in this mode. Jump directly to the
 *   parser-loop on reset
 *   F_CPU in BOOTICEMODE must be 7372800 Hz to be compatible
 *   with the org. JTAGICE-Firmware
 * WAIT-mode waits 1 sec for the defined character if nothing 
 *    is recived then the user prog is started.
 */
//#define START_SIMPLE
#define START_WAIT
//#define START_POWERSAVE
//#define START_BOOTICE

/* character to start the bootloader in mode START_WAIT */
#define START_WAIT_UARTCHAR 0xDF

/* wait 1s in START_WAIT mode (10ms steps) */
#define WAIT_VALUE	50
#define WAIT_MS		((F_CPU / 1000))

/*
 * enable/disable readout of fuse and lock-bits
 * (AVRPROG has to detect the AVR correctly by device-code
 * to show the correct information).
 */
//#define ENABLEREADFUSELOCK

 /* enable/disable write of lock-bits
 * WARNING: lock-bits can not be reseted by bootloader (as far as I know)
 * Only protection no unprotection, "chip erase" from bootloader only
 * clears the flash but does no real "chip erase" (this is not possible
 * with a bootloader as far as I know)
 * Keep this undefined!
 */
//#define WRITELOCKBITS

#define VERSION_HIGH '0'
#define VERSION_LOW  '8'

#define GET_LOCK_BITS           0x0001
#define GET_LOW_FUSE_BITS       0x0000
#define GET_HIGH_FUSE_BITS      0x0003
#define GET_EXTENDED_FUSE_BITS  0x0002

#include "bootldr.h"
#include "main.h"
#include "spm.h"
#include <stdint.h>
#include <avr/io.h>
#include <avr/wdt.h>
#include <avr/boot.h>
#include <avr/pgmspace.h>
#include <util/delay.h>
#include "chipdef.h"
#include <util/crc16.h>
#include "aes.h"
#include <avr/eeprom.h>

//-------------------------------------------------------------------------------------------------

typedef unsigned char byte;

//-------------------------------------------------------------------------------------------------

//=============================================================================
// Error Codes

#define ERROR_OK 0x11
#define ERROR_CRC 0x22

//---------------------------------------------------------------------
// Respone Codes
#define CMD_OK 1
#define CMD_FAILED 2
#define CMD_CRC_FAILED 3
#define CMD_NEXT 4
#define SIG_FAILED 5
#define CMD_BAD_HEADS 0x0A
#define BOOT_LOCKED 0x99

//=============================================================================
// Frame types

#define TYPE_EOF 0
#define TYPE_ERASE 1
#define TYPE_PREPARE 2
#define TYPE_DATA 3
#define TYPE_PROGRAM 4
#define TYPE_EEPROM 5
#define TYPE_LOCKBITS 6
#define TYPE_RESET 7

byte gBuffer[SPM_PAGESIZE];

#if KEY_COUNT > 0
byte initialVector[16] =
{
	(byte)(INITIALVECTOR_3 >> 24),
	(byte)(INITIALVECTOR_3 >> 16),
	(byte)(INITIALVECTOR_3 >> 8),
	(byte)(INITIALVECTOR_3 >> 0),
	(byte)(INITIALVECTOR_2 >> 24),
	(byte)(INITIALVECTOR_2 >> 16),
	(byte)(INITIALVECTOR_2 >> 8),
	(byte)(INITIALVECTOR_2 >> 0),
	(byte)(INITIALVECTOR_1 >> 24),
	(byte)(INITIALVECTOR_1 >> 16),
	(byte)(INITIALVECTOR_1 >> 8),
	(byte)(INITIALVECTOR_1 >> 0),
	(byte)(INITIALVECTOR_0 >> 24),
	(byte)(INITIALVECTOR_0 >> 16),
	(byte)(INITIALVECTOR_0 >> 8),
	(byte)(INITIALVECTOR_0 >> 0)
};
#endif // KEY_COUNT > 0

union {
	struct {
		byte m_rxBuffer[BUFFER_SIZE];	// Receive buffer
		byte m_pageBuffer[PAGE_SIZE];	// Page is assembled here before
							// getting programmed to flash mem
	} part1;

#if KEY_COUNT > 0
	struct {
		byte m_tempbuf[256]; // Temp buffer for aesInit.
	} part2;
#endif
} sharedbufs;

#define rxBuffer   sharedbufs.part1.m_rxBuffer
#define pageBuffer sharedbufs.part1.m_pageBuffer
#if KEY_COUNT > 0
	#define tempbuf    sharedbufs.part2.m_tempbuf
#endif


#if defined(BOOTLOADERHASNOVECTORS)
#warning "This Bootloader does not link interrupt vectors - see makefile"
/* make the linker happy - it wants to see __vector_default */
// void __vector_default(void) { ; }
void __vector_default(void) { ; }
#endif


static void sendchar(uint8_t data)
{
	while (!(UART_STATUS & (1<<UART_TXREADY)));
	UART_DATA = data;
}

static uint8_t recvchar(void)
{
	while (!(UART_STATUS & (1<<UART_RXREADY)));
	return UART_DATA;
}

static inline void eraseFlash(void)
{
	// erase only main section (bootloader protection)
	uint32_t addr = 0;
	while (APP_END > addr) {

		//spmErasePage(&addr);

		boot_page_erase(addr);		// Perform page erase
		boot_spm_busy_wait();		// Wait until the memory is erased.
		addr += SPM_PAGESIZE;
	}
	boot_rww_enable();
}

static inline void recvBuffer(pagebuf_t size)
{
	pagebuf_t cnt;
	uint8_t *tmp = gBuffer;

	for (cnt = 0; cnt < sizeof(gBuffer); cnt++)
		*tmp++ = (cnt < size) ? recvchar() : 0xFF;
}

static inline uint32_t writeFlashPage(uint32_t waddr, pagebuf_t size)
{
	uint32_t pagestart = (uint32_t)waddr;
	uint32_t baddr = pagestart;
	uint16_t data;
	uint8_t *tmp = gBuffer;

	do {
		data = *tmp++;
		data |= *tmp++ << 8;
		boot_page_fill(baddr, data);	// call asm routine.

		baddr += 2;			// Select next word in memory
		size -= 2;			// Reduce number of bytes to write by two
	} while (size);				// Loop until all bytes written

	boot_page_write(pagestart);
	boot_spm_busy_wait();
	boot_rww_enable();			// Re-enable the RWW section

	return baddr;
}

static inline uint16_t writeEEpromPage(uint16_t address, pagebuf_t size)
{
	uint8_t *tmp = gBuffer;

	do {
		spmEEWriteByte(address++, *tmp++);
		//address++;
		//tmp++;

		//EEARL = address;		// Setup EEPROM address
		//EEARH = (address >> 8);
		//EEDR = *tmp++;
		//address++;			// Select next byte

		//EECR |= (1<<EEMWE);		// Write data into EEPROM
		//EECR |= (1<<EEWE);
		//eeprom_busy_wait();

		size--;				// Decreas number of bytes to write
	} while (size);				// Loop until all bytes written

	return address;
}

static inline uint16_t readFlashPage(uint32_t waddr, pagebuf_t size)
{
	uint32_t baddr = waddr;
	uint16_t data;

	do {
#if defined(RAMPZ)
		data = pgm_read_word_far(baddr);
#else
		data = pgm_read_word_near(baddr);
#endif
		sendchar(data);			// send LSB
		sendchar((data >> 8));		// send MSB
		baddr += 2;			// Select next word in memory
		size -= 2;			// Subtract two bytes from number of bytes to read
	} while (size);				// Repeat until all block has been read

	return baddr;
}

static inline uint16_t readEEpromPage(uint16_t address, pagebuf_t size)
{
	do 
	{
		EEARL = address;		// Setup EEPROM address
		EEARH = (address >> 8);
		EECR |= (1<<EERE);		// Read EEPROM
		address++;			// Select next EEPROM byte

		sendchar(EEDR);			// Transmit EEPROM data to PC

		size--;				// Decrease number of bytes to read
	} while (size);				// Repeat until all block has been read

	return address;
}

static inline byte EEpromRead(uint16_t address)
{
	EEARL = address;		// Setup EEPROM address
	EEARH = (address >> 8);
	EECR |= (1<<EERE);	
	
	return EEDR;
}

static inline void Comm_header(void)
{ sendchar(0xAA); sendchar(0xA5); sendchar(0x1A); sendchar(0x70); }

static inline void Comm_footer(void)
{ sendchar(0x55); sendchar(0x5A); }

static void Comm_Response(byte resp)
{
	unsigned int crc=0;
	
	Comm_header();
	sendchar(0x00);
	sendchar(0x02);
	sendchar(0xAA);
	sendchar(resp);
	
	crc = _crc16_update(crc, 0xAA);
	crc = _crc16_update(crc, resp);
	sendchar((crc >> 8));
	sendchar(crc);
	
	Comm_footer();
}

#if defined(ENABLEREADFUSELOCK)
static uint8_t read_fuse_lock(uint16_t addr)
{
	uint8_t mode = (1<<BLBSET) | (1<<SPMEN);
	uint8_t retval;

	asm volatile
	(
		"movw r30, %3\n\t"		/* Z to addr */ \
		"sts %0, %2\n\t"		/* set mode in SPM_REG */ \
		"lpm\n\t"			/* load fuse/lock value into r0 */ \
		"mov %1,r0\n\t"			/* save return value */ \
		: "=m" (SPM_REG),
		  "=r" (retval)
		: "r" (mode),
		  "r" (addr)
		: "r30", "r31", "r0"
	);
	return retval;
}
#endif

static void send_boot(void)
{
	unsigned int crc = 0;
	unsigned int adr = 0;
	byte EE = 0;
	Comm_header();
	sendchar(0x00);
	sendchar(0x05);
	crc = _crc16_update(crc, 0x55);
	sendchar(0x55);

	sendchar(BOOTBYTE);
	crc = _crc16_update(crc, BOOTBYTE);

	adr = 0x000C;
	EE = EEpromRead(adr);
	sendchar(EE);
	crc = _crc16_update(crc, EE);
	adr = 0x000D;
	EE = EEpromRead(adr);
	sendchar(EE);
	crc = _crc16_update(crc, EE);
	adr = 0x000E;
	EE = EEpromRead(adr);
	sendchar(EE);
	crc = _crc16_update(crc, EE);
	
	sendchar((crc >> 8));
	sendchar(crc);
	Comm_footer();

	// Send Board number here...
}

static void (*jump_to_app)(void) = 0x0000;

inline int16_t recvcharW(void)
{
	uint16_t timeOut = 0;
	while (!(UCSR0A & (1<<RXC0)))
	{
		timeOut++;
		wdt_reset();
		if (timeOut >= WAIT_MS)
		{
			return -1;
		}
	}
	return UDR0;
}


#if HWVER == 31
void wdt_init(void) __attribute__((naked)) __attribute__((section(".init3")));

// Function Implementation
void wdt_init(void)
{
    MCUSR = 0;
    wdt_disable();

    return;
}
#endif

int main(void)
{
	wdt_disable();

	#if KEY_COUNT > 0
		unsigned char chainCipherBlock[16];		// Buffer for Cipher Block
												// Unchaining
	
		// Copy Initial Vector to start the cipher block unchaining
		{
			unsigned char *p = chainCipherBlock;
			unsigned char *q = initialVector;
			unsigned char n = 16;
		
			do
			{
				*p++ = *q++;
			}
			while (--n);
		}
		
		aesInit (tempbuf); // Init AES algorithm.
	#endif

	cbi(DDR_SCK, SCKNUM);
	DDR_SCK  &= ~(1<<SCKNUM);		// set as Input
	PORT_SCK |= (1<<SCKNUM);		// Enable pullup

	// Set baud rate
	UART_BAUD_HIGH = (UART_CALC_BAUDRATE(BAUDRATE)>>8) & 0xFF;
	UART_BAUD_LOW = (UART_CALC_BAUDRATE(BAUDRATE) & 0xFF);

#ifdef UART_DOUBLESPEED
	UART_STATUS = UART_DOUBLE;
#endif

	UART_CTRL = UART_CTRL_DATA;
	UART_CTRL2 = UART_CTRL2_DATA;

#if defined(START_WAIT)

	uint16_t cnt = 0;
	byte byCmd[10] = {0xAA, 0xA5, 0x70, 0x1A, 0x00, 0x00, 0x00, 0x00, 0x55, 0x5A};
	// AAA5701A00000000555A
	byte byCmdlen = 0;
	int16_t rcv;

	if (bit_is_clear(PIN_SCK, SCKNUM))
	{
		while (1) 
		{
			rcv = recvcharW();
			if (rcv != -1)
			{
				if (byCmd[byCmdlen] == rcv)
				{	
					byCmdlen++; 
				}
				else
				{
					PORT_SCK &= ~(1<<SCKNUM);
					jump_to_app();
				}
				if (byCmdlen == 10)
				{ 
					send_boot();
					break;
				}
			}
			else if (cnt++ >= WAIT_VALUE)
			{
				PORT_SCK &= ~(1<<SCKNUM);		// set to default
				jump_to_app();			// Jump to application sector
			}
		}
	}
	else
	{ 
		//Comm_Response(CMD_OK);
		jump_to_app();
	}		// Jump to application sector

#elif defined(START_BOOTICE)
#warning "BOOTICE mode - no startup-condition"

#else
#error "Select START_ condition for bootloader in main.c"
#endif

	// Enter the Flash Uploading mode, now erase flash and begin writing
	eraseFlash();
	if ((boot_lock_fuse_bits_get(GET_LOCK_BITS) & 0x0C) == 0x0c)
		Comm_Response(CMD_OK);
	else
	{
		 Comm_Response(BOOT_LOCKED);
		 jump_to_app();
	}
	
	unsigned int iSize = 0;

	uint16_t crc;
	frameindex_t frameSize;
	unsigned int crcEx;
	frameindex_t n;
	unsigned char ch;
	
	for(;;) {	
		// Get the frame size
		if ((recvchar() == 0xAA) && 
		    (recvchar() == 0xA5) &&
			(recvchar() == 0x70) &&
			(recvchar() == 0x1A))
		{
			frameSize =	((frameindex_t)recvchar() << 8) |	recvchar();
	
			// Receive a frame of data from communication interface and calculate
			// its CRC
			unsigned char *p = rxBuffer;
			crcEx = 0;
			n = frameSize;//-2;
	
			crc = 0;
			do
			{
				ch = recvchar();
	
				*p++ = ch;
				crc = _crc16_update(crc, ch);
			}
			while (--n);
			
			p += 2;
		
			crcEx = ((unsigned int)recvchar() << 8) | recvchar();
			
			if ((recvchar() == 0x55) && (recvchar() == 0x5A))
			{
				if (crcEx == crc)
				{
					unsigned char *p = rxBuffer;			
		
					// Decrypt 16 bytes, CRC-bytes are ignored.
					#if KEY_COUNT > 0
					{
						unsigned char *p = rxBuffer;
		
						//frameSize -= 2;
						do
						{
							//__watchdog_reset();
							aesDecrypt(p, chainCipherBlock);
							p += 16;
						}
						while (frameSize -= 16);
					}
					#endif // KEY_COUNT > 0
		
					// Check that the signature is correct
					if ((*p++ == (unsigned char)(SIGNATURE >> 24)) &&
						(*p++ == (unsigned char)(SIGNATURE >> 16)) &&
						(*p++ == (unsigned char)(SIGNATURE >>  8)) &&
						(*p++ == (unsigned char)(SIGNATURE >>  0)))
					{
						unsigned long address;
						unsigned char bits;
						frameindex_t size;
						unsigned char type;
		
						// Continue parsing the frames until the 'End Of Frame' is
						// received
						if ((type = *p++) != TYPE_EOF)
						//while ((type = *p++) != TYPE_EOF)
						{
							//	__watchdog_reset();
		
							// Receive Lock Bits (used in TYPE_LOCKBITS)
							bits = *p++;
		
							// Destination Address (note: 'bits' and 'address' overlap)
		
							address = ((unsigned long)bits << 16);
							address = ((unsigned long)address | ((unsigned long)*p++ << 8));
							address = ((unsigned long)address | *p++);
		
							size = ((unsigned int)*p++ << 8);
							size |= *p++;
		
							switch (type)
							{
								// Chunk of data is written into the page buffer
								case TYPE_DATA:
								{
									unsigned char *r = gBuffer; //&gBuffer[(unsigned long)(address) & 0xffff];
									iSize = size;
									do
									{
										*r++ = *p++;
									}
									while (--size);
									writeFlashPage(address, iSize);
									//iAdr = writeFlashPage(iAdr, iSize);
									Comm_Response(CMD_OK);
									break;
								}
		
								// Program page buffer into flash page
								case TYPE_PROGRAM:
								{
									//readFlashPage(address, iSize);
									//writeFlashPage(address, iSize);
									//readFlashPage(address, iSize);
									Comm_Response(CMD_OK);
									break;
								}
								
								case TYPE_PREPARE:
								{
									unsigned char *r = gBuffer; //&gBuffer[(unsigned long)(address) & 0xffff];
									iSize = size;
									do
									{
										*r++ = *p++;
									}
									while (--size);
									writeFlashPage(address, iSize);
									//iAdr = writeFlashPage(iAdr, iSize);
									Comm_Response(CMD_OK);
									break;
								} 
		
								// Write a chunk of data into EEPROM
								case TYPE_EEPROM:
									eeprom_write_block((const void*)p, (void*)(uint16_t)(address & 0xFFFF), size);
									Comm_Response(CMD_OK);
									break;
		
								// Write Lock Bits
								case TYPE_LOCKBITS:
									//spmWriteLockBits(bits);
									break;
		
								// Reset (let the Watchdog Timer expire)!
								case TYPE_RESET:
									wdt_enable(WDTO_15MS);
									for(;;)
									break;
									
								// Nonsense frames, one byte long
								default:
									p -= 5;
									break;
							}
						}
					}
					else
					{ Comm_Response(SIG_FAILED); }
				}
				else
				{ Comm_Response(CMD_CRC_FAILED); }
			}
			else
			{ Comm_Response(CMD_BAD_HEADS);	}
		}
		else
		{ Comm_Response(CMD_BAD_HEADS);	}
	}
	
	return 0;
}
